from typing import List
import os
import numpy as np
import torch
import time
import argparse
import logging


class NamespaceDict(dict, argparse.Namespace):
    """
    A dictionary that supports both attribute-style and dictionary-style access.

    Allows accessing, setting, and deleting keys as attributes or dictionary items.

    Example:
        d = NamespaceDict(a=1, b=2)
        print(d.a)  # 1
        print(d['b'])  # 2
        d.c = 3
        print(d['c'])  # 3
        del d.a
        print(hasattr(d, 'a'))  # False
    """

    def __init__(self, *args, **kwargs):
        """Initialize the dict and set attributes from the provided items."""
        dict.__init__(self, *args, **kwargs)
        for key, value in self.items():
            self.__setattr__(key, value)

    def __getitem__(self, key):
        return dict.__getitem__(self, key)

    def __getattr__(self, key):
        # Called when accessing non-existent attributes; fallback to dict
        try:
            return dict.__getitem__(self, key)
        except KeyError:
            raise AttributeError(f"'NamespaceDict' object has no attribute '{key}'")

    def __delitem__(self, key):
        """Allow deletion by both dict-style and attribute-style."""
        dict.__delitem__(self, key)
        super().__delattr__(key)

    def __setattr__(self, name, value):
        if isinstance(value, (list, tuple)):
            value = type(value)(
                self.__class__(x) if isinstance(x, dict) else x for x in value
            )
        elif isinstance(value, dict) and not isinstance(value, NamespaceDict):
            value = NamespaceDict(value)
        super(NamespaceDict, self).__setattr__(name, value)
        super(NamespaceDict, self).__setitem__(name, value)

    __setitem__ = __setattr__
    __delattr__ = __delitem__

    def update(self, *args, **kwargs):
        dict.update(self, *args, **kwargs)
        for key, value in dict(*args, **kwargs).items():
            self.__setattr__(key, value)

    def pop(self, key, default=None):
        """Pop the key from the dictionary and attributes."""
        val = dict.pop(self, key, default)
        super().__delattr__(key)
        return val

    def __eq__(self, other):
        return vars(self) == vars(other)

    def __contains__(self, key):
        return key in self.__dict__

    def __repr__(self):
        type_name = type(self).__name__
        arg_strings = []
        star_args = {}
        for arg in self._get_args():
            arg_strings.append(repr(arg))
        for name, value in self._get_kwargs():
            if name.isidentifier():
                arg_strings.append("%s=%r" % (name, value))
            else:
                star_args[name] = value
        if star_args:
            arg_strings.append("**%s" % repr(star_args))
        return "%s(%s)" % (type_name, ", ".join(arg_strings))

    __str__ = __repr__


class NamespaceDictBP1(dict, argparse.Namespace):
    """
    A dictionary that supports both attribute-style and dictionary-style access.

    Allows accessing, setting, and deleting keys as attributes or dictionary items.

    Example:
        d = NamespaceDict(a=1, b=2)
        print(d.a)  # 1
        print(d['b'])  # 2
        d.c = 3
        print(d['c'])  # 3
        del d.a
        print(hasattr(d, 'a'))  # False
    """

    def __init__(self, *args, **kwargs):
        """Initialize the dict and set attributes from the provided items."""
        dict.__init__(self, *args, **kwargs)
        for key, value in self.items():
            super().__setattr__(key, value)

    def __getitem__(self, key):
        return dict.__getitem__(self, key)

    def __getattr__(self, key):
        # Called when accessing non-existent attributes; fallback to dict
        try:
            return dict.__getitem__(self, key)
        except KeyError:
            raise AttributeError(f"'NamespaceDict' object has no attribute '{key}'")

    def __setitem__(self, key, value):
        """Allow dict-style setting and set attributes."""
        dict.__setitem__(self, key, value)
        super().__setattr__(key, value)

    def __delitem__(self, key):
        """Allow deletion by both dict-style and attribute-style."""
        dict.__delitem__(self, key)
        super().__delattr__(key)

    __setattr__ = __setitem__
    __delattr__ = __delitem__

    def update(self, *args, **kwargs):
        """Update the dictionary and attributes."""
        dict.update(self, *args, **kwargs)
        for key, value in kwargs.items():
            super().__setattr__(key, value)

    def pop(self, key, default=None):
        """Pop the key from the dictionary and attributes."""
        val = dict.pop(self, key, default)
        super().__delattr__(key)
        return val

    def __eq__(self, other):
        return vars(self) == vars(other)

    def __contains__(self, key):
        return key in self.__dict__

    def __repr__(self):
        type_name = type(self).__name__
        arg_strings = []
        star_args = {}
        for arg in self._get_args():
            arg_strings.append(repr(arg))
        for name, value in self._get_kwargs():
            if name.isidentifier():
                arg_strings.append("%s=%r" % (name, value))
            else:
                star_args[name] = value
        if star_args:
            arg_strings.append("**%s" % repr(star_args))
        return "%s(%s)" % (type_name, ", ".join(arg_strings))

    __str__ = __repr__


class NamespaceDictbp(argparse.Namespace, dict):
    def __init__(self, *args, **kwargs):
        # Initialize dict with given arguments
        dict.__init__(self, *args, **kwargs)
        # Set all keys as attributes
        for key, value in self.items():
            super().__setattr__(key, value)

    def __getitem__(self, key):
        # Allow dict-style access
        return dict.__getitem__(self, key)

    def __setitem__(self, key, value):
        # Allow dict-style setting and also set attributes
        dict.__setitem__(self, key, value)
        super().__setattr__(key, value)

    def __setattr__(self, key, value):
        # Set both dict key and attribute when using dot notation
        dict.__setitem__(self, key, value)
        super().__setattr__(key, value)

    def __delitem__(self, key):
        # Allow deletion by both dict-style and attribute-style
        dict.__delitem__(self, key)
        super().__delattr__(key)

    def __delattr__(self, key):
        # Allow attribute-style deletion
        try:
            dict.__delitem__(self, key)
        except KeyError:
            pass
        super().__delattr__(key)

    def __getattr__(self, key):
        # Called when accessing non-existent attributes; fallback to dict
        try:
            return dict.__getitem__(self, key)
        except KeyError:
            raise AttributeError(f"'NamespaceDict' object has no attribute '{key}'")


def time_execution(func):
    def wrapper(*args, **kwargs):
        start_time = time.time()  # Start timer
        result = func(*args, **kwargs)  # Call the original function
        end_time = time.time()  # End timer
        elapsed_time = end_time - start_time
        print(f"Time taken for {func.__name__:25s}: {elapsed_time:.2f} seconds")
        return result

    return wrapper


def count_parameters(model, trainable=False):
    if trainable:
        return sum(p.numel() for p in model.parameters() if p.requires_grad)
    return sum(p.numel() for p in model.parameters())


def tensor2numpy(x, dtype=np.float32):
    return (
        x.cpu().data.numpy().astype(dtype)
        if x.is_cuda
        else x.data.numpy().astype(dtype)
    )


def target2onehot(targets, n_classes):
    onehot = torch.zeros(targets.shape[0], n_classes).to(targets.device)
    onehot.scatter_(dim=1, index=targets.long().view(-1, 1), value=1.0)
    return onehot


def makedirs(path):
    if not os.path.exists(path):
        os.makedirs(path)


def accuracy_dil(y_pred, y_true, nb_old, inc_cls, cls_num=0):
    assert len(y_pred) == len(y_true), "Data length error."
    all_acc = {}
    all_acc["total"] = np.around(
        (y_pred % cls_num == y_true % cls_num).sum() * 100 / len(y_true),
        decimals=2,
    )

    # Grouped accuracy
    for class_id in range(0, np.max(y_true), inc_cls):
        idxes = np.where(
            np.logical_and(y_true >= class_id, y_true < class_id + inc_cls)
        )[0]
        label = "{}-{}".format(
            str(class_id).rjust(2, "0"),
            str(class_id + inc_cls - 1).rjust(2, "0"),
        )
        all_acc[label] = np.around(
            ((y_pred[idxes] % cls_num) == (y_true[idxes] % cls_num)).sum()
            * 100
            / len(idxes),
            decimals=2,
        )

    # Old accuracy
    idxes = np.where(y_true < nb_old)[0]
    all_acc["old"] = (
        0
        if len(idxes) == 0
        else np.around(
            ((y_pred[idxes] % cls_num) == (y_true[idxes] % cls_num)).sum()
            * 100
            / len(idxes),
            decimals=2,
        )
    )

    # New accuracy
    idxes = np.where(y_true >= nb_old)[0]
    all_acc["new"] = np.around(
        ((y_pred[idxes] % cls_num) == (y_true[idxes] % cls_num)).sum()
        * 100
        / len(idxes),
        decimals=2,
    )

    return all_acc


def accuracy_cil(y_pred, y_true, nb_old, inc_cls):
    assert len(y_pred) == len(y_true), "Data length error."
    all_acc = {}
    all_acc["total"] = np.around(
        (y_pred == y_true).sum() * 100 / len(y_true), decimals=2
    )

    # FIXME: This is a temporary solution to load the ViT models for fine-tuning
    inc_cls = inc_cls if inc_cls > 0 else np.max(y_true)
    # Grouped accuracy
    for class_id in range(0, np.max(y_true), inc_cls):
        idxes = np.where(
            np.logical_and(y_true >= class_id, y_true < class_id + inc_cls)
        )[0]
        label = "{}-{}".format(
            str(class_id).rjust(2, "0"),
            str(class_id + inc_cls - 1).rjust(2, "0"),
        )
        all_acc[label] = np.around(
            (y_pred[idxes] == y_true[idxes]).sum() * 100 / len(idxes),
            decimals=2,
        )

    # Old accuracy
    idxes = np.where(y_true < nb_old)[0]

    all_acc["old"] = (
        0
        if len(idxes) == 0
        else np.around(
            (y_pred[idxes] == y_true[idxes]).sum() * 100 / len(idxes),
            decimals=2,
        )
    )

    # New accuracy
    idxes = np.where(y_true >= nb_old)[0]
    all_acc["new"] = np.around(
        (y_pred[idxes] == y_true[idxes]).sum() * 100 / len(idxes),
        decimals=2,
    )

    return all_acc


def accuracy(y_pred, y_true, nb_old, inc_cls):
    assert len(y_pred) == len(y_true), "Data length error."
    all_acc = {}
    all_acc["total"] = np.around(
        (y_pred == y_true).sum() * 100 / len(y_true), decimals=2
    )

    # FIXME: This is a temporary solution to load the ViT models for fine-tuning
    inc_cls = inc_cls if inc_cls > 0 else np.max(y_true)
    # Grouped accuracy
    for class_id in range(0, np.max(y_true), inc_cls):
        idxes = np.where(
            np.logical_and(y_true >= class_id, y_true < class_id + inc_cls)
        )[0]
        label = "{}-{}".format(
            str(class_id).rjust(2, "0"),
            str(class_id + inc_cls - 1).rjust(2, "0"),
        )
        all_acc[label] = np.around(
            (y_pred[idxes] == y_true[idxes]).sum() * 100 / len(idxes),
            decimals=2,
        )

    # Old accuracy
    idxes = np.where(y_true < nb_old)[0]

    all_acc["old"] = (
        0
        if len(idxes) == 0
        else np.around(
            (y_pred[idxes] == y_true[idxes]).sum() * 100 / len(idxes),
            decimals=2,
        )
    )

    # New accuracy
    idxes = np.where(y_true >= nb_old)[0]
    all_acc["new"] = np.around(
        (y_pred[idxes] == y_true[idxes]).sum() * 100 / len(idxes),
        decimals=2,
    )

    return all_acc


def split_images_labels(imgs):
    # split trainset.imgs in ImageFolder
    images = []
    labels = []
    for item in imgs:
        images.append(item[0])
        labels.append(item[1])

    return np.array(images), np.array(labels)


class AccMeter:
    """Computes and stores average and current values for an arbitrary number of metrics (e.g., losses)."""

    def __init__(self, names, fmt=".4f", tail=""):
        """Initialize with names of metrics."""
        self.names = (
            names  # List of metric names (e.g., ['loss1', 'loss2', 'accuracy'])
        )
        self.fmt = fmt
        self.tail = tail
        self.reset()

    def add(self, names: List):
        """Add new metrics to the meter."""
        for name in names:
            if name in self.names:
                continue
            self.names.append(name)
            self.avg[name] = 0.0
            self.sum[name] = 0.0
            self.count[name] = 0.0

    def reset(self):
        """Reset all meters."""
        self.avg = {name: 0.0 for name in self.names}
        self.sum = {name: 0.0 for name in self.names}
        self.count = {name: 0.0 for name in self.names}

    def update(self, dict_vals):
        """Update all metrics with new values."""
        for name, val in dict_vals.items():
            self.sum[name] += val["correct"]
            self.count[name] += val["total"]
            self.avg[name] = (self.sum[name] / self.count[name]) * 100

    def __str__(self):
        """Return formatted string for displaying all metrics."""
        metrics_str = []
        for name in self.names:
            fmtstr = f"{name:10s}: {{avg:{self.fmt}}}" + self.tail
            metrics_str.append(fmtstr.format(avg=self.avg[name], name=name))
        return "  ".join(metrics_str)


class LossMeter(AccMeter):
    def update(self, dict_vals):
        for name, val in dict_vals.items():
            self.sum[name] += val
            self.count[name] += 1
            self.avg[name] = self.sum[name] / self.count[name]


class MeterProcessor:
    """Handles the display of metric progress across epochs."""

    def __init__(self, name, epochs, meters, prefix=""):
        self.name, self.epoch_fmt, self.meters, self.prefix = (
            name,
            self.__epoch_fmtstr(epochs),
            meters,
            prefix,
        )

    def display(self, epoch, *args, **kwargs):
        """Display metrics for the current epoch."""
        entries = [f"{self.prefix}{self.epoch_fmt.format(epoch)}"] + [
            str(meter) for meter in self.meters
        ]
        logging.info("  ".join(entries))

    def __epoch_fmtstr(self, epochs):
        """Generate format string for displaying epoch information."""
        num_digits = len(str(int(epochs)))
        return f"[{{:0{num_digits}d}}/{epochs}]"


class TrainMP(MeterProcessor):
    """Extended MeterProcessor for displaying training-specific metrics (including multiple losses)."""

    def __init__(self, name, epochs, meters, prefix=""):
        super().__init__(name, epochs, meters, prefix)
        self.timer = None

    def reset_timer(self):
        """Start the timer."""
        self.timer = time.time()

    def log_timer(self):
        """Stop the timer and return elapsed time."""
        if self.timer is not None:
            elapsed_time = time.time() - self.timer
            return elapsed_time
        return None

    def display(self, epoch, lr: List):
        """Display training metrics along with timer and learning rate."""
        elapsed_time = self.log_timer()
        timer_str = f"[{elapsed_time:.2f} seconds]" if elapsed_time else "[No Timer]"
        lr_str = f"{[f'{i:.3e}' for i in lr]}\n"
        entries = [
            f"{self.prefix}{self.epoch_fmt.format(epoch)} {timer_str} {lr_str}"
        ] + [str(meter) for meter in self.meters]
        logging.info("  ".join(entries))
